import pandas as pd
from cvzone.FaceMeshModule import FaceMeshDetector
import cv2
import torch
import pyautogui as pag
from torch.utils.tensorboard import SummaryWriter
import keyboard
import warnings
warnings.filterwarnings("ignore")

img_name_index = 0
bool_add = True
writer = SummaryWriter("logs")
leftEyeUpPoint = []
leftEyeDownPoint = []
leftEyeLeftPoint = []
leftEyeRightPoint = []
rightEyeUpPoint = []
rightEyeDownPoint = []
rightEyeLeftPoint = []
rightEyeRightPoint = []
# 初始化网络摄像头
# '2' 表示与计算机连接的第三台相机，'0' 通常指内置网络摄像头
cap = cv2.VideoCapture(0)

# 初始化 FaceMeshDetector 对象
# staticMode: 若为真，则仅检测一次，否则每帧都检测
# maxFaces: 最大可检测人脸数量
# minDetectionCon: 检测置信度阈值
# minTrackCon: 跟踪置信度阈值
detector = FaceMeshDetector(staticMode=False, maxFaces=1, minDetectionCon=0.5, minTrackCon=0.5)

# 人眼关键点所在的索引
idList = [22, 23, 24, 26, 110, 130, 157, 158, 159, 160, 161, 243]  #左眼

# 开始循环以持续获取摄像头帧
while True:
    # 从网络摄像头读取当前帧
    # success: 帧是否成功捕获的布尔值
    # img: 当前帧
    success, img = cap.read()

    # 在图像中找到面部网格
    # img: 如设置 draw=True 则更新了面部网格的图像
    # faces: 检测到的面部信息
    img, faces = detector.findFaceMesh(img, draw=False)
    # 如果检测到面部
    if faces:
        # 遍历每个检测到的面部
        for face in faces:
            # font = cv2.FONT_HERSHEY_SIMPLEX
            # for (index, value) in enumerate(face):
            #     cv2.putText(img, str(index), value, font, 0.2, (255, 230, 0), 1, cv2.LINE_AA)
            # 获取眼睛特定点
            leftEyeUpPoint = face[159]
            leftEyeDownPoint = face[23]
            leftEyeLeftPoint = face[130]
            leftEyeRightPoint = face[243]
            rightEyeUpPoint = face[386]
            rightEyeDownPoint = face[253]
            rightEyeLeftPoint = face[463]
            rightEyeRightPoint = face[466]

            # 绘制眼部周围的特征点
            demo1 = [leftEyeUpPoint, leftEyeDownPoint, leftEyeLeftPoint, leftEyeRightPoint]
            demo2 = [rightEyeUpPoint, rightEyeDownPoint, rightEyeLeftPoint, rightEyeRightPoint]
            # for i in demo1:
            #     cv2.circle(img, i, 1, (0, 0, 0), -1)
            # for i in demo2:
            #     cv2.circle(img, i, 1, (0, 0, 0), -1)

            # 框选出眼睛的矩形图
            aa1x = leftEyeLeftPoint[0] - 10
            aa1y = min(leftEyeUpPoint[1], rightEyeUpPoint[1]) - 10
            aa2x = rightEyeRightPoint[0] + 10
            aa2y = max(rightEyeDownPoint[1], leftEyeDownPoint[1]) + 10
            aa1 = [aa1x, aa1y]
            aa2 = [aa2x, aa2y]

            # cv2.rectangle(img,aa1, aa2, (255, 0, 0), 2)

            def cropped_eyes():
                x = aa1x
                y = aa1y
                w = abs(aa1x - aa2x)
                h = abs(aa1y - aa2y)
                return x, y, w, h

            x, y, w, h = cropped_eyes()
            cropped_screen = img[y:y + h, x:x + w]


            def get_image():
                img_begin = torch.from_numpy(cropped_screen)
                image = torch.tensor(img_begin).permute(2, 0, 1)
                image = image.unsqueeze(0)
                return image.float() / 127. - 1.


            def add_data():
                global img_name_index
                global bool_add
                if bool_add:
                    label_data = pd.read_csv('./data/train/label/labels.csv')

                    screenWidth, screenHeight = pag.size()
                    x, y = pag.position()
                    # 返回鼠标的坐标
                    print('屏幕： (%s,%s)  鼠标坐标 : (%s, %s)' % (screenWidth, screenHeight ,x, y))

                    # 使用concat添加新行
                    new_row = pd.DataFrame([[img_name_index,(x/screenWidth)*2-1, (y/screenHeight)*2-1]], columns=['ID','x', 'y'])  # 创建一个只包含新行的DataFrame
                    label_data = pd.concat([label_data, new_row], ignore_index=True)  # 将新行添加到原始DataFrame

                    label_data.to_csv('./data/train/label/labels.csv', index=False)
                    cv2.imwrite("./data/train/img/" + str(img_name_index) + ".png", cropped_screen)

                    img_name_index += 1
                    tensor_img = get_image().numpy()

                    print(tensor_img.shape)
                    writer.add_images("eyes", tensor_img)
                    bool_add = False


            # 按下空格收集数据集
            keyboard.add_hotkey(" ", add_data, suppress=False)
            bool_add = True

    # 在名为 'Image' 的窗口中显示图像
    cv2.imshow("Image", img)

    # 等待 1 毫秒检查用户输入，保持窗口打开状态
    cv2.waitKey(1)

# 打开tensorboard
#tensorboard --logdir=logs --port=6001